Skip to content

Gate solve_tril dot precision on whether TF32 is faster, not whether it exists - #1205

Open
arbi-dev wants to merge 1 commit into
fla-org:mainfrom
arbi-dev:fix/solve-tril-precision-not-faster-tf32
Open

Gate solve_tril dot precision on whether TF32 is faster, not whether it exists#1205
arbi-dev wants to merge 1 commit into
fla-org:mainfrom
arbi-dev:fix/solve-tril-precision-not-faster-tf32

Conversation

@arbi-dev

@arbi-dev arbi-dev commented Aug 30, 2026

Copy link
Copy Markdown

Summary

IS_TF32_SUPPORTED answers "is TF32 available" — its own comment says as
much ("Nvidia Ampere or newer"). The three fused solve_tril sites use it to
answer a different question: "is TF32 worth taking". On the datacenter dies
those two answers agree. On the GeForce/Ada dies they do not.

The case

Each site inverts a block-triangular WY factor with
tl.dot(..., input_precision=SOLVE_TRIL_DOT_PRECISION) on operands that are
fp32 accumulators (tl.zeros([BC, BC], dtype=tl.float32)) — so the setting is
live, not inert.

Part FP32 TF32 (dense) Worth taking?
A100 (sm_80) 19.5 TFLOPS 156 TFLOPS yes, 8x
RTX 4090 (sm_89) 82.6 TFLOPS 82.6 TFLOPS no, 1.0x

On sm_89 the gate spends 13 mantissa bits (23 → 10) and buys no throughput — on
a triangular inverse, the step of the chunk that amplifies error most.

This is the direction the library already leans, for this same operation

  • fla/ops/utils/solve_tril.py:19-22 — the non-fused implementation of the
    same op defaults to ieee and only admits tf32 into its autotune list
    under IS_TMA_SUPPORTED (sm_90+):
    FLA_TRIL_PRECISION = os.environ.get('FLA_TRIL_PRECISION', 'ieee')
    DOT_PRECISION_AUTOTUNE_LIST = ["ieee"] if not IS_TMA_SUPPORTED else list({"ieee", FLA_TRIL_PRECISION})
  • fla/ops/utils/matmul.py:192,230 — declines tf32 whenever the operands
    are fp32:
    allow_tf32 = False if a.dtype == torch.float32 else True

So the three fused copies look like the outliers rather than the intent. This
PR is less "change the policy" than "apply the existing one to the sites keyed
off availability".

Changes

  • IS_TF32_FASTER_THAN_FP32 in fla/utils/_device.py, beside
    IS_TF32_SUPPORTED, keyed off an explicit TF32_NOT_FASTER_CAPABILITIES
    set. Only (8, 9) is listed, so no unlisted part changes behaviour. RTX
    30xx (8.6) and Blackwell GeForce (12.0) are believed to be in the same
    position but are left out — I would rather under-claim than assert specs I
    cannot source. Happy to add them if you know them.
  • All three fused sites converted together:
    gated_delta_rule/chunk_fwd.py, gdn2/chunk_intra.py, kda/chunk_intra.py
    carry byte-identical copies of the branch, so fixing one would leave two.

Capability-tuple comparison follows the existing IS_NVIDIA_SM120 pattern in
the same file. No new env var — FLA_TRIL_PRECISION already exists for the
non-fused path, and extending it to cover these is your API call, not mine.

Risk

No behaviour change on any datacenter part. The default moves only on
sm_89; any capability not in the set keeps today's answer. The direction on
sm_89 is toward precision, so it needs no accuracy case — the throughput
claim carries the argument, and that is from published specs rather than a
benchmark here.

Testing

Verified all three modules resolve from the shared predicate and agree, and
that the predicate answers False for (8, 9) and True for (8, 0) /
(9, 0). I have no sm_80/sm_90 card to hand, so the datacenter path rests on
its default being unchanged rather than on a measurement. Glad to add a unit
test for the predicate if you want one in tests/.

Context

Found while auditing a GDN serving path that runs Qwen3.5/3.6 hybrids on sm_89.
We are carrying this as a local override (arbicity/arbi-serve#1739) and would
much rather drop it than keep it.

…it exists

`IS_TF32_SUPPORTED` answers "is TF32 available" — its own comment says so
("Nvidia Ampere or newer"). The three fused solve_tril sites use it to answer a
different question, "is TF32 worth taking". On the datacenter dies those agree;
on the GeForce/Ada dies they do not.

Each site inverts a block-triangular WY factor with tl.dot(...,
input_precision=SOLVE_TRIL_DOT_PRECISION) on operands that are fp32
accumulators (tl.zeros([BC, BC], dtype=tl.float32)), so the setting is live.
On A100 selecting tf32 there is clearly right: fp32 19.5 vs tf32 156 TFLOPS.
On sm_89 it is not — an RTX 4090 runs fp32 and tf32 alike at 82.6 TFLOPS — so
the gate spends 13 mantissa bits and buys no throughput, on the step of the
chunk that amplifies error most.

The library already leans the other way for this exact operation elsewhere:
ops/utils/solve_tril.py defaults FLA_TRIL_PRECISION to 'ieee' and only puts
tf32 in its autotune list under IS_TMA_SUPPORTED (sm_90+). And
ops/utils/matmul.py declines tf32 whenever the operands are fp32
(allow_tf32 = False if a.dtype == torch.float32 else True). The three fused
copies are the outliers rather than the intent.

Adds IS_TF32_FASTER_THAN_FP32 next to IS_TF32_SUPPORTED, keyed off an explicit
TF32_NOT_FASTER_CAPABILITIES set. Only (8, 9) is listed, so no unlisted part
changes behaviour; 8.6 and 12.0 are believed to be in the same position and are
left out until someone can confirm them. Capability-tuple comparison follows
the existing IS_NVIDIA_SM120 pattern in the same file.

All three sites are converted together — gated_delta_rule/chunk_fwd.py,
gdn2/chunk_intra.py and kda/chunk_intra.py carry byte-identical copies of the
branch, so fixing one would leave two.

No behaviour change on any datacenter part. Verified all three modules resolve
from the shared predicate and agree; I have no sm_80/sm_90 card to hand, so the
datacenter path rests on the default being unchanged rather than on a
measurement.
@arbi-dev
arbi-dev force-pushed the fix/solve-tril-precision-not-faster-tf32 branch from 4c8f510 to 840dbb4 Compare August 30, 2026 15:44
@arbi-dev arbi-dev changed the title Pick solve_tril dot precision on whether TF32 is faster, not whether it exists Gate solve_tril dot precision on whether TF32 is faster, not whether it exists Aug 30, 2026
@zhiyuan1i zhiyuan1i added the bug Something isn't working label Aug 31, 2026

@zhiyuan1i zhiyuan1i left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Style issues:

  1. P1: TF32_NOT_FASTER_CAPABILITIES is exported from fla/utils/__init__.py but has no external callers — unjustified new public symbol. Remove from __init__.py.

  2. P2: IS_TF32_FASTER_THAN_FP32 is not in _register_aliases() but the constant it replaces (IS_TF32_SUPPORTED) is. Either add the alias or confirm downstream doesn't need it.

  3. P2: 9-line comment in _device.py:153-161 — first sentence restates what the identifier already says, A100 vs 4090 TFLOPS comparison can be one sentence. Keep only "sm_89 fp32/tf32 same 82.6 TFLOPS" and "unlisted models unchanged", ~3-4 lines.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants